#include <getopt.h>
#include <iostream>
#include <fstream>
#include <iomanip>
#include <vector>
#include <gmssl/pbkdf2.h>
#include <gmssl/rand.h>
#include <gmssl/sm9.h>
#include <gmssl/sm4.h>
using namespace std;
option long_opt[]={
    {"password",1,0,'p'},
    {"verify",0,0,'v'},
    {"genkey",0,0,'g'},
    {"help",0,0,'h'},
    {"sign",0,0,'n'},
    {"file",1,0,'f'},
    {"enc",0,0,'e'},
    {"dec",0,0,'d'},
    {"str",1,0,'s'},
    {"alg",1,0,'a'},
    {"out",1,0,'o'},
    {0,0,0,0}
};
void usage(){
    cout<<"Show uasge: --help/-help/-h\n";
    cout<<"Enc: --enc/-enc/-c\n";
    cout<<"Dec: --dec/-dec/-d\n";
    cout<<"Sign: --sign/-sign/-n\n";
    cout<<"Verify: --verify/-verify/-v\n";
    cout<<"Set in filename: --file/-file/-f filename\n";
    cout<<"Set out filename: --out/-out/-o filename\n";
    cout<<"Set plaintext/ciphertext: --str/-str/-s string\n";
    cout<<"Select algorithm: --alg/-alg/-a algorithm\n";
    cout<<"Generate a key: --genkey/-genkey/-g";
    cout<<"Set sm4 generate key password: --password/-password/-p";
    cout<<"-------------------------------------------------\n";
    cout<<"Algorithms: sm2,sm4,sm9\n";
    exit(0);
}
vector<uint8_t> randomData(size_t len){
    vector<uint8_t> res(len);
    rand_bytes(res.data(),len);
    return res;
}
vector<uint8_t> sm2_encrypt2(const SM2_KEY* key,const vector<uint8_t>& in){
    size_t len=in.size()+SM2_MAX_CIPHERTEXT_SIZE;
    vector<uint8_t> res(len);
    sm2_encrypt(key,in.data(),in.size(),res.data(),&len);
    res.resize(len);
    return res;
}
vector<uint8_t> sm2_decrypt2(const SM2_KEY* key,const vector<uint8_t>& in){
    size_t len=in.size();
    vector<uint8_t> res(len);
    sm2_decrypt(key,in.data(),in.size(),res.data(),&len);
    res.resize(len);
    return res;
}
vector<uint8_t> sm9_encrypt2(const SM9_ENC_MASTER_KEY* mpk,const vector<uint8_t>& in,const string& id){
    size_t len=in.size()+SM9_MAX_CIPHERTEXT_SIZE;
    vector<uint8_t> res(len);
    sm9_encrypt(mpk,id.c_str(),id.size(),in.data(),in.size(),res.data(),&len);
    res.resize(len);
    return res;
}
vector<uint8_t> sm9_decrypt2(const SM9_ENC_KEY* key,const vector<uint8_t>& in,const string& id){
    size_t len=in.size();
    vector<uint8_t> res(len);
    sm9_decrypt(key,id.data(),id.size(),in.data(),in.size(),res.data(),&len);
    res.resize(len);
    return res;
}
vector<uint8_t> sm4_encrypt2(
    const vector<uint8_t>& key,
    const vector<uint8_t>& iv,
    const vector<uint8_t>& aad,
    const vector<uint8_t>& plaintext,
    vector<uint8_t>& tag) {
    SM4_KEY sm4_key;
    sm4_set_encrypt_key(&sm4_key, key.data());
    vector<uint8_t> ciphertext(plaintext.size());
    tag.resize(SM4_GCM_MAX_TAG_SIZE);
    sm4_gcm_encrypt(
        &sm4_key,iv.data(),iv.size(),
        aad.data(),aad.size(),
        plaintext.data(),plaintext.size(),
        ciphertext.data(),
        tag.size(),tag.data()
    );
    return ciphertext;
}
vector<uint8_t> sm4_decrypt2(
    const vector<uint8_t>& key,
    const vector<uint8_t>& iv,
    const vector<uint8_t>& aad,
    const vector<uint8_t>& ciphertext,
    const vector<uint8_t>& tag) {
    SM4_KEY sm4_key;
    sm4_set_encrypt_key(&sm4_key, key.data());
    vector<uint8_t> plaintext(ciphertext.size());
    if(sm4_gcm_decrypt(
        &sm4_key,iv.data(),iv.size(),
        aad.data(),aad.size(),
        ciphertext.data(),ciphertext.size(),
        tag.data(),tag.size(),
        plaintext.data())!=1){
        throw runtime_error("GCM标签验证失败");
    }
    return plaintext;
}
void printHex(const vector<uint8_t>& in){
    cout<<hex<<setfill('0');
    for(auto& i:in) cout<<setw(2)<<(short)i;
    cout<<dec<<setfill(' ')<<endl;
}
vector<uint8_t> genSm4Key(string pass){
    vector<uint8_t> res,salt=randomData(32);
    pbkdf2_hmac_sm3_genkey(pass.c_str(),pass.size(),salt.data(),32,10000,16,res.data());
    return res;
}
int main(int argc,char** argv) {
    int doType=-1,type=-1;
    string filename,message,password,ofilename;
    char ch;
    int index;
    while((ch=getopt_long_only(argc,argv,"hedvngp:f:a:s:o:",long_opt,&index))!=-1){
        switch(ch){
        case 'a':
            if(optarg=="sm2") type=0;
            else if(optarg=="sm4") type=1;
            else if(optarg=="sm9") type=2;
            else{
                usage();
                return 0;
            }
            break;
        case 'o':ofilename=optarg;break;
        case 's':message=optarg;break;
        case 'p':password=optarg;break;
        case 'f':filename=optarg;break;
        case 'e':doType=0;break;
        case 'd':doType=1;break;
        case 'n':doType=2;break;
        case 'v':doType=3;break;
        case 'g':doType=4;break;
        case 'h':
        default: usage();
        }
    }
    if(doType==-1||type==-1||(filename==""&&message=="")||
      (filename!=""&&message!="")) usage();
    if(filename!=""){
        ifstream in(filename);
        if(!in.is_open()) usage();
        string line;
        while(getline(in,line)) message+=line+'\n';
        in.close();
    }
    return 0;
}